import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


class Quadrotor:
    def __init__(self):
        # 物理参数
        self.m = 0.5  # 质量 (kg)
        self.g = 9.81  # 重力加速度 (m/s²)
        self.L = 0.25  # 旋翼到重心的距离 (m)
        self.Ixx = 0.027  # 绕x轴转动惯量 (kg·m²)
        self.Iyy = 0.027  # 绕y轴转动惯量 (kg·m²)
        self.dt = 0.1  # 时间步长 (s)

        # 状态变量：[x, y, z, phi, theta, psi, vx, vy, vz, p, q, r]
        # 这里简化为6维状态（假设角速度p, q, r为0）
        self.n_states = 6
        self.n_controls = 4  # [F, M_x, M_y, M_z]

        # 状态空间模型
        self.A = np.eye(self.n_states)
        self.A[:3, 3:6] = self.dt * np.eye(3)  # 速度影响位置

        # 输入矩阵（简化模型）
        self.B = np.zeros((self.n_states, self.n_controls))
        self.B[3, 0] = self.dt / self.m  # 总力F影响垂直加速度
        self.B[4, 1] = self.dt * self.L / self.m  # 力矩M_y影响X方向加速度
        self.B[5, 2] = self.dt * self.L / self.m  # 力矩M_x影响Y方向加速度

    def step(self, state, control):
        return self.A @ state + self.B @ control


class TrajectoryGenerator:
    def __init__(self, total_time):
        self.total_time = total_time

    def get_desired_state(self, t):
        # 生成螺旋轨迹（三维）
        radius = 2.0
        omega = 0.2
        z_amplitude = 0.5
        x_d = radius * np.cos(omega * t)
        y_d = radius * np.sin(omega * t)
        z_d = z_amplitude * np.sin(omega * t)
        vx_d = -radius * omega * np.sin(omega * t)
        vy_d = radius * omega * np.cos(omega * t)
        vz_d = z_amplitude * omega * np.cos(omega * t)
        return np.array([x_d, y_d, z_d, vx_d, vy_d, vz_d])


class MPCController:
    def __init__(self, quadrotor, N, Q, R):
        self.quadrotor = quadrotor
        self.N = N  # 预测时域
        self.Q = Q  # 状态权重矩阵（6x6）
        self.R = R  # 控制输入权重矩阵（4x4）

    def control(self, state, desired_states):
        n_states = self.quadrotor.n_states
        n_controls = self.quadrotor.n_controls
        x = cp.Variable((n_states, self.N + 1))
        u = cp.Variable((n_controls, self.N))

        cost = 0
        constraints = []

        # 初始状态约束
        constraints.append(x[:, 0] == state)

        for k in range(self.N):
            # 状态转移约束
            constraints.append(
                x[:, k + 1] ==
                self.quadrotor.A @ x[:, k] +
                self.quadrotor.B @ u[:, k]
            )

            # 代价函数
            cost += cp.quad_form(x[:, k] - desired_states[:, k], self.Q)
            cost += cp.quad_form(u[:, k], self.R)

        # 终端代价
        cost += cp.quad_form(x[:, -1] - desired_states[:, -1], self.Q)

        prob = cp.Problem(cp.Minimize(cost), constraints)
        prob.solve(solver=cp.OSQP)

        return u[:, 0].value


def simulate():
    quadrotor = Quadrotor()
    traj_gen = TrajectoryGenerator(total_time=20)
    N = 5  # 预测时域
    Q = np.diag([1.0, 1.0, 1.0, 0.1, 0.1, 0.1])  # 状态权重（位置权重高）
    R = np.diag([0.1, 0.01, 0.01, 0.01])  # 输入权重（总力权重较高）
    mpc = MPCController(quadrotor, N, Q, R)

    # 初始化状态（位置0，速度0）
    state = np.zeros(quadrotor.n_states)
    actual_trajectory = []
    desired_trajectory = []

    total_time = 20
    num_steps = int(total_time / quadrotor.dt)

    for i in range(num_steps):
        t = i * quadrotor.dt
        # 生成未来N步的期望状态
        desired_states = np.zeros((quadrotor.n_states, N + 1))
        for k in range(N + 1):
            desired_states[:, k] = traj_gen.get_desired_state(t + k * quadrotor.dt)

        # 计算控制输入
        control = mpc.control(state, desired_states)

        # 更新状态
        state = quadrotor.step(state, control)

        # 存储轨迹
        actual_trajectory.append(state[:3])
        desired_trajectory.append(desired_states[:3, 0])

    # 绘制三维轨迹
    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(111, projection='3d')
    actual = np.array(actual_trajectory)
    desired = np.array(desired_trajectory)

    ax.plot(desired[:, 0], desired[:, 1], desired[:, 2], 'b--', label='Desired')
    ax.plot(actual[:, 0], actual[:, 1], actual[:, 2], 'r-', label='Actual')
    ax.set_xlabel('X (m)'), ax.set_ylabel('Y (m)'), ax.set_zlabel('Z (m)')
    ax.legend()
    ax.set_title('3D Trajectory Tracking')

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    simulate()